import sys
import cudf
import cugraph
import socket
import json

client_address = ('127.0.0.1', 31500)
server_address = ('127.0.0.1', 31501)
server_s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
server_s.bind(server_address)

G = cugraph.Graph()

if __name__== "__main__":
    graph_file = "src_dst.csv"
    if len(sys.argv) >= 2:
        graph_file = sys.argv[1]

    distance = None
    bfs_src = -1

    while True:
        # cmd = input("graph>>>")
        print("graph>>>")
        json_string, addr = server_s.recvfrom(2048)
        cmd = json.loads(json_string)
        cmd_list = cmd.split(" ")
        print(cmd_list)
        f = open("out.txt", "w")
        if cmd_list[0] == "load":
            print("loading graph...", file=f)
            gdf = cudf.read_csv("src_dst.csv", names=["src", "dst"], dtype=["int32", "int32"])
            G.clear()
            G.from_cudf_edgelist(gdf, source='src', destination='dst', renumber=False)
            print("graph load success", file=f)
        elif cmd_list[0] == "pagerank":
            sort_pr = cugraph.pagerank(G).sort_values('pagerank', ascending=False)
            topk = int(cmd_list[1])
            for i in range(topk): #len(sort_pr)
                print("vertex " + str(sort_pr['vertex'].iloc[i]), file=f)#+ " PageRank is " + str(sort_pr['pagerank'].iloc[i]))
        elif cmd_list[0] == "sssp":
            src = int(cmd_list[1])
            distance = cugraph.sssp(G,src)
            print(distance, file=f)
        elif cmd_list[0] == "bfs":
            try:
                src = int(cmd_list[1])
                distance = cugraph.bfs(G,src)
                print(distance, file=f)
            except:
                pass
        elif cmd_list[0] == "path":
            try:
                src = int(cmd_list[1])
                dst = int(cmd_list[2])
                if bfs_src != src:
                    bfs_src = src
                    distance = cugraph.bfs(G,src)
                print("distance:", distance['distance'].iloc[dst], file=f)
                path=[dst]
                pre = distance['predecessor'].iloc[dst]
                while pre != -1:
                    path.append(pre)
                    pre = distance['predecessor'].iloc[pre]
                path.reverse()
                print("PATH:", "->".join([str(i) for i in path]), file=f)
            except:
                pass
        elif cmd_list[0] == "node":
            try:
                node = int(cmd_list[1])
                print(G.has_node(node), file=f)
            except:
                pass
        elif cmd_list[0] == "edge":
            src = int(cmd_list[1])
            dst = int(cmd_list[2])
            print(G.has_edge(src, dst), file=f)
        elif cmd_list[0] == "info":
            print("Graph has {} nodes".format(G.number_of_nodes()), file=f)
            print("Graph has {} edges".format(G.number_of_edges()), file=f)
        elif cmd_list[0] == "neighbor":
            try:
                node = int(cmd_list[1])
                print(G.neighbors(node), file=f)
            except:
                pass
        elif cmd_list[0] == "hops":
            try:
                node = int(cmd_list[1])
                hops = int(cmd_list[2])
                if hops == 0:
                    print(node, file=f)
                else:
                    result = [[node]]
                    for i in range(hops):
                        for x in result[i]:
                            result.append(list(G.neighbors(x)))
                    result_list = []
                    for n in result[1:]:
                        result_list.extend(n)
                    result_set = set(result_list)
                    print(result_set, file=f)
            except:
                pass
        else:
            print("error cmd")
            break
        f.close()
        server_s.sendto(bytes(json.dumps("SUCCESS"),'UTF-8'), client_address)

    server_s.sendto(bytes(json.dumps("SUCCESS"),'UTF-8'), client_address)
    server_s.close()